"""
添加数据库索引以提升查询性能
"""
import sys
import os

# 添加父目录到路径
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from sqlalchemy import create_index, Index, text
from app.database import engine, SessionLocal
from app.models import Document, Note, Category, KnowledgeBase
import logging

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)


def add_indexes():
    """添加数据库索引"""
    db = SessionLocal()
    
    try:
        logger.info("=" * 60)
        logger.info("开始添加数据库索引...")
        logger.info("=" * 60)
        
        # 使用原生 SQL 添加索引（如果不存在）
        indexes = [
            # Documents 表索引
            ("idx_documents_user_id", "documents", "user_id"),
            ("idx_documents_category_id", "documents", "category_id"),
            ("idx_documents_kb_id", "documents", "knowledge_base_id"),
            ("idx_documents_created_at", "documents", "created_at"),
            ("idx_documents_title", "documents", "title"),
            
            # Notes 表索引
            ("idx_notes_user_id", "notes", "user_id"),
            ("idx_notes_kb_id", "notes", "knowledge_base_id"),
            ("idx_notes_created_at", "notes", "created_at"),
            ("idx_notes_title", "notes", "title"),
            
            # Categories 表索引
            ("idx_categories_user_id", "categories", "user_id"),
            ("idx_categories_parent_id", "categories", "parent_id"),
            
            # Knowledge Bases 表索引
            ("idx_knowledge_bases_user_id", "knowledge_bases", "user_id"),
            ("idx_knowledge_bases_visibility", "knowledge_bases", "visibility"),
        ]
        
        for idx_name, table_name, column_name in indexes:
            try:
                # 检查索引是否已存在
                check_sql = text(f"""
                    SELECT 1 FROM pg_indexes 
                    WHERE indexname = '{idx_name}'
                """)
                result = db.execute(check_sql).fetchone()
                
                if result:
                    logger.info(f"⏭️  索引 {idx_name} 已存在，跳过")
                    continue
                
                # 创建索引
                create_sql = text(f"""
                    CREATE INDEX {idx_name} 
                    ON {table_name} ({column_name})
                """)
                db.execute(create_sql)
                db.commit()
                logger.info(f"✅ 创建索引: {idx_name} on {table_name}({column_name})")
                
            except Exception as e:
                logger.warning(f"⚠️  索引 {idx_name} 创建失败: {str(e)}")
                db.rollback()
                continue
        
        logger.info("=" * 60)
        logger.info("✅ 数据库索引添加完成！")
        logger.info("=" * 60)
        
    except Exception as e:
        logger.error(f"❌ 添加索引失败: {str(e)}")
        import traceback
        traceback.print_exc()
    finally:
        db.close()


if __name__ == "__main__":
    add_indexes()

